# Copyright (c) 2020 cgsdfc
# 
# This software is released under the MIT License.
# https://opensource.org/licenses/MIT

import ast
from pprint import pprint
from sys import version_info
from typing import List
from pathlib import Path
from io import StringIO
import re
import os
import mmap
import astlib.asdl as asdl
import pycparser
from pycparser.c_ast import (FileAST, Typedef, Decl, Enum, Struct, TypeDecl,
                             Union, Enumerator, EnumeratorList)
from tempfile import mkstemp
from itertools import filterfalse

MISSING_TYPES = (
    'asdl_seq',
    'asdl_int_seq',
    'object',
    'string',
    'identifier',
    'constant',
    'singleton',
    'bytes',
    'PyArena',
    'PyObject',
)

PP_DIRECTIVE_RE = re.compile(r'^#.*$')


def ParsePythonAsdlHeader(filename):
    # (str) ->FileAST
    missig_decls = ['typedef int {};'.format(t) for t in MISSING_TYPES]
    data = open(filename).read()
    code = missig_decls + data.split('\n')[1:]
    code = list(filterfalse(PP_DIRECTIVE_RE.match, code))
    code = missig_decls + code

    _, fname = mkstemp(suffix='.c')
    with open(fname, 'w') as f:
        f.write('\n'.join(code))
    result = pycparser.parse_file(filename=fname)
    os.unlink(fname)
    return result


def IterElements(node):
    for _, ch in node.children():  # type: None, Decl
        # Enum like operator and unaryop.
        if isinstance(ch, Typedef):
            if isinstance(ch.type, TypeDecl) and isinstance(
                    ch.type.type, Enum):
                yield ch.name.replace('_ty', ''), ch.type.type

        if isinstance(ch, Decl) and isinstance(ch.type, Enum):
            enum = ch.type  # type: Enum
            # Enum like _mod_kind
            if enum.name.endswith('_kind'):
                for e in enum.values.enumerators:  # type: Enumerator
                    yield e.name.replace('_kind', ''), None

        # Struct like expr, stmt and mod.
        if isinstance(ch, Decl) and isinstance(ch.type, Struct):
            struct = ch.type  # type: Struct
            yield struct.name.lstrip('_'), struct

            # Struct like BinOp and UnaryOp.
            for d in struct.decls:  # type: Decl
                if d.name == 'v' and isinstance(d.type.type, Union):
                    union = d.type.type  # type: Union
                    for decl in union.decls:
                        yield decl.name, decl.type.type


class LocalNs:
    def __init__(self, fields):
        self._fields = fields

    def __repr__(self):
        return f'<LocalNs {self._fields}>'


def HasEndLocationBySysVersion():
    # TODO: find the exact version.
    return version_info.major == 3 and version_info.minor >= 8


def ExtractNamespaceFromBuiltinAstModule():
    module = {}
    for name in dir(ast):
        obj = getattr(ast, name)
        if isinstance(obj, type) and issubclass(obj, ast.AST):
            module[name] = obj
    return module, HasEndLocationBySysVersion()


def ExtractNamespaceFromHeaderAst(node: FileAST):
    global_ns = {}
    has_end_location = False

    for name, ch in IterElements(node):  # type: str, Union[Struct, Enum]
        if isinstance(ch, Enum):
            global_ns[name] = None
            for e in ch.values.enumerators:  # type: Enumerator
                global_ns[e.name] = None
        elif isinstance(ch, Struct):
            d = ch.decls[0]
            has_local_ns = ch.name is None or not (
                d.name == 'kind' and isinstance(d.type.type, Enum))

            if has_local_ns:
                fields = set(d.name for d in ch.decls)
                # Remove lineno, col_offset pair.
                if 'lineno' in fields and 'col_offset' in fields:
                    fields.remove('lineno')
                    fields.remove('col_offset')
                if 'end_lineno' in fields and 'end_col_offset' in fields:
                    fields.remove('end_lineno')
                    fields.remove('end_col_offset')
                    has_end_location = True

                global_ns[name] = LocalNs(frozenset(fields))
            else:
                global_ns[name] = None
        elif ch is None:
            # A constructor type without a struct def.
            global_ns.setdefault(name)
        else:
            raise TypeError('unexpected ch: {}'.format(ch))
    return global_ns, has_end_location


# HAS_END_LOCATION = HasEndLocation()
ASDL_TEMPLATE = Path(__file__).with_name('Python.asdl')


class MaterializingError(Exception):
    pass


def RemoveNones(seq):
    return list(filter(None, seq))


class MaterializingVisitor(asdl.VisitorBase):
    """
    This visitor deletes nodes of a `asdl.Module` to match the definitions
    in the builtin `ast` module. It also ensures a one-to-one mapping from
    the asdl nodes to the ast module's classes.
    """
    def __init__(self, header=None):
        super().__init__()
        if header:
            # This set makes sure the reduced namespace of asdl is the same
            # as the module.
            self.asdl_names = set()
            node = ParsePythonAsdlHeader(header)
            # This mapping makes sure asdl_names in a subset of module.
            module, has_end_location = ExtractNamespaceFromHeaderAst(node)
        else:
            self.asdl_names = {'AST'}
            module, has_end_location = ExtractNamespaceFromBuiltinAstModule()

        self.module = module
        self.has_end_location = has_end_location

    def visitModule(self, mod: asdl.Module):
        dfns_copy = list(mod.dfns)
        for i, dfn in enumerate(dfns_copy):  # type: int, asdl.Type
            name = dfn.name
            if name not in self.module:
                mod.dfns[i] = None
                del mod.types[name]
            else:
                self.visit(dfn)
        mod.dfns = RemoveNones(mod.dfns)
        mod.has_end_location = self.has_end_location

        # Final check that the two namespaces is equal.
        module_names = set(self.module.keys())
        if module_names != self.asdl_names:
            raise MaterializingError(
                'The asdl global namespace is smaller than that of ast.\n'
                'module - asdl: {}\n'
                'asdl - module: {}'.format(module_names - self.asdl_names,
                                           self.asdl_names - module_names))

    def visitType(self, type: asdl.Type):
        assert type.name in self.module
        self.asdl_names.add(type.name)
        if isinstance(type.value, asdl.Product):
            # Set the product's name to that of type.
            type.value.name = type.name
        self.visit(type.value)

    def _visitAttributes(self, attributes: List[asdl.Field]):
        if attributes and not self.has_end_location:
            assert len(attributes) == 4
            del attributes[2:]  # Remove end_lineno and end_col_offset

    def visitSum(self, sum_: asdl.Sum):
        types_copy = list(sum_.types)  # Get a copy.
        for i, type in enumerate(types_copy):  # type: int, asdl.Constructor
            name = type.name
            if name not in self.module:
                # del sum_.types[i]
                sum_.types[i] = None
            else:
                self.visit(type)
        sum_.types = RemoveNones(sum_.types)
        self._visitAttributes(sum_.attributes)

    def _visitFields(self, node, name=None):
        # No LocalNs for this symbol.
        name = name or node.name
        if self.module[name] is None:
            return

        fields = node.fields
        fields_copy = tuple(f.name for f in node.fields)
        ast_fields = self.module[name]._fields
        ast_fields_ns = frozenset(ast_fields)

        for i, field_name in enumerate(fields_copy):  # type: int, asdl.Field
            if field_name not in ast_fields_ns:
                fields[i] = None
        node.fields = RemoveNones(fields)

        asdl_fields = tuple(f.name for f in node.fields)
        asdl_fields_ns = set(asdl_fields)
        if ast_fields_ns != asdl_fields_ns:
            raise MaterializingError(
                'The asdl local namespace is smaller than that of ast.\n'
                'name: {}\n'
                'ast_fields: {}\n'
                'asdl_fields: {}'.format(name, ast_fields, asdl_fields))

    def visitConstructor(self, ctor: asdl.Constructor):
        assert ctor.name in self.module
        self.asdl_names.add(ctor.name)
        self._visitFields(ctor)

    def visitProduct(self, product: asdl.Product):
        # Product's name is the name of the Type.
        assert hasattr(product, 'name')
        self._visitAttributes(product.attributes)
        self._visitFields(product)



def DumpAsdl(module: asdl.AST, file=None):
    vis = UnparsingVisitor(file)
    vis.visit(module)
    if file is None:
        return vis.ostream.getvalue()


def LoadAsdl(filename):
    return asdl.parse(filename)


def MaterializePythonAsdl(header=None):
    """
    Load a template Python.asdl and materialize it using the ast module.
    """
    module = asdl.parse(str(ASDL_TEMPLATE))
    vis = MaterializingVisitor(header)
    vis.visit(module)
    return module
